import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
# from utils import tools

import os
import cv2
import PIL.Image as Image
import  numpy as np
from torch.utils.data import DataLoader
import torch.optim as optim
from network.resnet100 import KitModel
import torch.nn.functional as F
import torchvision.transforms as transforms
from face_modules.mtcnn import *
import torch
# from apex import amp

# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
from sync_batchnorm import convert_model, patch_replication_callback
from network.sexual import *



test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

detector = MTCNN()
arcface_weight_path = os.path.join('./Arcface_100.pth')
# arcface = KitModel(arcface_weight_path).to(device)  # used to generate face embed feature vector
arcface = sexdiscrimination(arcface_weight_path).to(device)
arcface.eval()
# Xt_path = 'H:\\shipinjechen\\picture\\00006763.jpg'
Xt_path = "/home/zsy/桌面/sex/imgs/00000001.jpg"
Xs_path = '/home/zsy/桌面/sex/imgs/00000002.jpg'
Xt_raw = cv2.imread(Xt_path)
Xs_raw = cv2.imread(Xs_path)
# cv2.imshow("Xt_raw",Xt_raw )
# Xt = detector.align(Image.fromarray(Xt_raw[:, :, ::-1]), crop_size=(256, 256))
# Xs = detector.align(Image.fromarray(Xs_raw[:, :, ::-1]), crop_size=(256, 256))
Xt = detector.align(Image.fromarray(Xt_raw[:, :, ::-1]), crop_size=(224, 224))
Xs = detector.align(Image.fromarray(Xs_raw[:, :, ::-1]), crop_size=(224, 224))
Xt = Image.fromarray(np.array(Xt))
Xs = Image.fromarray(np.array(Xs))
Xt = test_transform(Xt)
Xs = test_transform(Xs)
Xt = Xt.unsqueeze(0).cuda()
Xs = Xs.unsqueeze(0).cuda()

embedt = arcface(F.interpolate(Xt[:, :, 19:237, 19:237], (112, 112), mode='bilinear', align_corners=True))
embeds = arcface(F.interpolate(Xs[:, :, 19:237, 19:237], (112, 112), mode='bilinear', align_corners=True))
print(embedt)
xt_r=embedt.detach().cpu().numpy()
xs_r=embeds.detach().cpu().numpy()

l=np.linalg.norm(xt_r-xs_r)
print(l)
arcface.train()